#include <algorithm>
#include <chrono>
#include <cstring>
#include <iostream>
#include <mutex>
#include <netinet/in.h>
#include <sys/select.h>
#include <sys/socket.h>
#include <thread>
#include <unistd.h>
#include <unordered_map>
#include <vector>

#define BUFFER_SIZE 1024 // 接收缓冲区大小
#define SELECT_TIMEOUT 3 // select超时时间（秒）

using namespace std;
using port_t = uint16_t;
using sock_t = int;

// 数据处理函数类型定义
using select_handler = void (*)(sock_t fd, const char *buffer, ssize_t len);

// 监听器结构体
struct Listener {
	sock_t sockfd;          // 套接字文件描述符
	port_t port;            // 绑定端口号
	select_handler handler; // 数据处理回调函数

	Listener(sock_t s = -1, port_t p = 0) : sockfd(s), port(p), handler(nullptr) {}
};

// UCP协议处理函数
void ucp_handler(sock_t fd, const char *buffer, ssize_t len) { std::cout << "[UCP处理] 来自端口 " << fd << " 的数据: " << buffer << std::endl; }

// 网络上下文类（线程安全）
class Context {
private:
public:
	std::unordered_map<sock_t, Listener *> sock2listener; // socket到监听器的映射
	std::unordered_map<port_t, Listener *> port2listener; // 端口到监听器的映射
	std::mutex listeners_mutex;                           // 访问互斥锁

	// 添加监听器（生命周期管理由调用方负责）
	void addListener(Listener *listener) {
		std::lock_guard<std::mutex> lock(listeners_mutex);
		std::cout << "[上下文] 添加监听器 - socket: " << listener->sockfd << " 端口: " << listener->port << std::endl;
		sock2listener[listener->sockfd] = listener;
		port2listener[listener->port]   = listener;
	}

	// 通过socket获取监听器
	Listener *getListenerByFD(sock_t fd) {
		std::lock_guard<std::mutex> lock(listeners_mutex);
		auto it = sock2listener.find(fd);
		return (it != sock2listener.end()) ? it->second : nullptr;
	}

	// 通过端口获取监听器
	Listener *getListenerByPort(port_t port) {
		std::lock_guard<std::mutex> lock(listeners_mutex);
		auto it = port2listener.find(port);
		return (it != port2listener.end()) ? it->second : nullptr;
	}

	// 移除监听器
	void removeListener(Listener *listener) {
		std::lock_guard<std::mutex> lock(listeners_mutex);
		std::cout << "[上下文] 移除监听器 - socket: " << listener->sockfd << " 端口: " << listener->port << std::endl;
		sock2listener.erase(listener->sockfd);
		port2listener.erase(listener->port);
		delete listener; // 释放内存
	}
};

// 网络选择器类
class Selector {
private:
	std::atomic<bool> running{false}; // 运行状态标志
	std::thread worker_thread;        // 工作线程
	Context &context;                 // 网络上下文引用

	// 工作线程函数
	void worker() {
		while (running) {
			std::cout << "\n[选择器] 开始扫描服务..." << std::endl;

			// 获取当前所有监听器
			vector<Listener *> listeners;
			{
				std::lock_guard<std::mutex> lock(context.listeners_mutex);
				for (auto &pair: context.sock2listener) {
					listeners.push_back(pair.second);
				}
			}

			// 准备select参数
			fd_set readfds;
			FD_ZERO(&readfds);
			int max_fd = 0;

			for (auto listener: listeners) {
				FD_SET(listener->sockfd, &readfds);
				max_fd = max(max_fd, listener->sockfd);
			}

			// 执行select
			timeval timeout{.tv_sec = SELECT_TIMEOUT, .tv_usec = 0};
			int ret = select(max_fd + 1, &readfds, nullptr, nullptr, &timeout);

			if (ret < 0) {
				std::cerr << "[选择器] select错误: " << strerror(errno) << std::endl;
				continue;
			} else if (ret == 0) {
				std::cout << "[选择器] 超时，无活动连接" << std::endl;
				continue;
			}

			// 处理就绪socket
			for (auto listener: listeners) {
				if (FD_ISSET(listener->sockfd, &readfds)) {
					char buffer[BUFFER_SIZE] = {0};
					sockaddr_in client_addr{};
					socklen_t addr_len = sizeof(client_addr);

					// 接收数据
					ssize_t len = recvfrom(listener->sockfd, buffer, BUFFER_SIZE, 0, (sockaddr *) &client_addr, &addr_len);

					if (len > 0) {
						std::cout << "[接收器] 收到来自 " << ntohl(client_addr.sin_addr.s_addr) << ":" << ntohs(client_addr.sin_port) << " 的数据: " << buffer << std::endl;

						// 调用处理函数
						if (listener->handler) {
							listener->handler(listener->sockfd, buffer, len);
						}
					}
				}
			}
		}
		std::cout << "[选择器] 工作线程退出" << std::endl;
	}

public:
	Selector(Context &ctx) : context(ctx) {
		running       = true;
		worker_thread = thread(&Selector::worker, this);
		std::cout << "[选择器] 初始化完成" << std::endl;
	}

	~Selector() {
		running = false;
		if (worker_thread.joinable()) {
			worker_thread.join();
		}
		std::cout << "[选择器] 已销毁" << std::endl;
	}

	void start() {
		if (!running) {
			running       = true;
			worker_thread = thread(&Selector::worker, this);
			std::cout << "[选择器] 已重新启动" << std::endl;
		}
	}
};

// 创建监听器工厂函数
Listener *createListener(port_t port = 0) {
	int sockfd = socket(AF_INET, SOCK_DGRAM, 0);
	if (sockfd < 0) {
		std::cerr << "[监听器] 创建socket失败: " << strerror(errno) << std::endl;
		return nullptr;
	}

	sockaddr_in server_addr{};
	server_addr.sin_family      = AF_INET;
	server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
	server_addr.sin_port        = htons(port);

	if (::bind(sockfd, (sockaddr *) &server_addr, sizeof(server_addr)) == -1) {
		std::cerr << "[监听器] 绑定端口失败: " << strerror(errno) << std::endl;
		close(sockfd);
		return nullptr;
	}

	// 获取实际绑定端口
	socklen_t len = sizeof(server_addr);
	getsockname(sockfd, (sockaddr *) &server_addr, &len);
	port_t actual_port = ntohs(server_addr.sin_port);

	std::cout << "[监听器] 创建成功 - socket: " << sockfd << " 绑定端口: " << actual_port << std::endl;
	return new Listener(sockfd, actual_port);
}


// g++ -std=c++11 -pthread -o selector.exe selector.cpp
int main() {
	Context context;
	vector<Listener *> listeners;

	// 创建两个监听器
	for (int i = 0; i < 2; ++i) {
		Listener *listener = createListener();
		if (listener) {
			listener->handler = ucp_handler; // 正确设置处理函数
			context.addListener(listener);
			listeners.push_back(listener);
		}
	}

	// 创建并启动选择器
	Selector selector(context);
	selector.start();

	// 发送测试消息
	for (size_t i = 0; i < 10; ++i) {
		std::cout << "\n[测试] 第 " << i + 1 << " 次消息发送" << std::endl;

		// 获取当前所有监听器
		vector<Listener *> current_listeners;
		{
			std::lock_guard<std::mutex> lock(context.listeners_mutex);
			for (auto &pair: context.sock2listener) {
				current_listeners.push_back(pair.second);
			}
		}

		// 向每个监听器发送消息
		for (auto listener: current_listeners) {
			string message = "Hello, Selector! " + to_string(i + 1);
			sendto(listener->sockfd, message.c_str(), message.size(), 0, nullptr, 0); // UDP广播发送
			std::cout << "[发送] 向端口 " << listener->port << " 发送: " << message << std::endl;
		}

		this_thread::sleep_for(chrono::seconds(1));
	}

	// 清理资源
	std::cout << "\n[清理] 开始释放资源..." << std::endl;
	for (auto listener: listeners) {
		context.removeListener(listener);
	}
	context.sock2listener.clear();
	context.port2listener.clear();
	std::cout << "[清理] 资源释放完成" << std::endl;

	return 0;
}
